#这是非线性方程组的迭代解法（牛顿法）
import numpy as np
from sympy import *
import pandas as pd
import warnings
import sys

warnings.filterwarnings("ignore")#忽略警告以避免无解报错
class Newton():
    def __init__(self, equations=[], variable=[], initial_value=[], e=None, K=None, pot=11):
        """
        这是牛顿法求解非线性方程组。
        :param equations: 方程组，以列表新式传递，内部为标准等式组。如：
                           x1**2+exp(x2)-10……省略“=0”
        :param variable:变量标志，如:x1,x2……
        :param initial_value:初始迭代值，与变量标志位置对应，默认全为0
        :param e:停止的相对误差限
        :param K:最大迭代次数
        :param pot:计算结果小数点位数，默认11，最大16位
        """
        np.set_printoptions(precision=pot)
        self.run = True
        self.variable = variable
        self.equations = equations
        self.e = e
        self.K = K
        if initial_value == []:
            self.x = np.zeros((len(self.variable), 1))
        else:
            self.x = np.array([initial_value]).T
            if self.x.shape[0] != len(self.variable):
                print("初始值输入有误！")
                self.run = False
        if len(self.variable) != len(self.equations):
            print("未知量与方程数目不等，无法求解。")
            self.run = False
        if self.variable==[] or self.equations==[] or self.e==None or self.K==None:
            self.run = False
        if self.run==True:
            self.to_run()

    def forward(self):
        Jf = np.zeros((len(self.variable), len(self.variable)))
        f = np.zeros((len(self.variable), 1))
        ans = np.zeros((len(self.variable), len(self.variable)))
        ans = pd.DataFrame(ans)
        for i in range(len(self.variable)):
            for j in range(len(self.variable)):
                ans.iloc[i, j] = diff(self.equations[i], self.variable[j])
                Jf[i][j] = diff(self.equations[i], self.variable[j]).subs({self.variable[k]: self.x[k][0] for k in range(len(self.variable))})
        try:
            for i in range(len(self.variable)):
                f[i] = self.equations[i].subs({self.variable[k]: self.x[k][0] for k in range(len(self.variable))})
        except:
            print("方程组中有错误字符，请检查！")
            sys.exit(1)
        self.ans = ans
        self.derta_x = np.linalg.inv(Jf) @ (-f)
        self.x = self.x + self.derta_x

    def to_run(self):
        try:
            for i in range(self.K):
                self.forward()
                if (self.e > abs(self.derta_x) / abs(self.x)).all():
                    break
        except:
            for i in range(self.K):
                self.forward()
                if (self.e > abs(self.derta_x) / abs(self.x)).all():
                    break
            print("迭代不收敛！请更换初始值试试！")
            sys.exit(1)
        print("求解结果：\n", self.x, "\n", "**"*20)
        return self.x

    def save_data(self, path):
        """
        保存计算数据的函数。
        :param path:保存路径及文件名（无文件类型后缀）。
        :return:
        """
        if self.run == False:
            print("保存失败，请检查输入！\n")
        else:
            eq = []
            for i in range(len(self.equations)):
                eq.append(str(self.equations[i]) + " = 0")
            for j in range(len(eq)):
                eq[j] = eq[j].replace("**", "^")
                eq[j] = eq[j].replace("log", "ln")
                eq[j] = eq[j].replace("log10", "log")
                eq[j] = eq[j].replace("exp", "e^")
                eq[j] = eq[j].replace("sin", "sin")
                eq[j] = eq[j].replace("cos", "cos")
                eq[j] = eq[j].replace("tan", "tan")
                eq[j] = eq[j].replace("arcsin", "arcsin")
                eq[j] = eq[j].replace("arccos", "arccos")
                eq[j] = eq[j].replace("arctan", "arctan")
            data = pd.DataFrame(data=[eq, self.variable, self.x], index=["方程组", "变量", "值"]).T
            self.ans.columns = ["梯度矩阵"] * len(self.variable)
            data = pd.concat((data, self.ans), axis=1)
            data.to_csv(path + "（ISONE）.csv", index=False)

def define_variables(n, s, start):
    """
    定义变量。直接定义n个以s为基本标记的变量，如s1 s2 s3……
    :param n: 变量数目 int型
    :param s: 变量标记 str型
    :param start:0 or 1,起始编号 s0 or s1
    :return:
    """
    if (type(n) is not int) or (type(s) is not str) or (start not in [0, 1]):
        print("变量定义部分输入有误！")
        sys.exit(1)
    for i in range(n):
        globals()[f'{s}{i + start}'] = symbols(f'{s}{i + start}', real=True) #定义变量

def one_stage_Newton(n, s, start, equations=[], variable=[], initial_value=[], e=0.01, K=1000, pot=11, save=True):
    """
    直接一次性实现牛顿非线性方程组迭代求解。
    :param n: 变量数目 int型
    :param s: 变量标记 str型
    :param start: 0 or 1,起始编号 s0 or s1
    :param equations: 方程组，以列表新式传递，内部为标准等式组。如：
                           “x1**2+exp(x2)-10”……省略“=0”，str型
    :param variable:变量标志，如:"x1",'x2'……，str型
    :param initial_value:初始迭代值，与变量标志位置对应，默认全为0
    :param e:停止的相对误差限,默认0.01
    :param K:最大迭代次数
    :param pot:计算结果小数点位数，默认11，最大16位
    :param save:结果保存标志，默认True，保存
    :return: 无
    """
    define_variables(n, s, start)
    for i in range(len(equations)):
        equations[i] = sympify(equations[i])
    for i in range(len(variable)):
        variable[i] = symbols(variable[i])
    in_one_stage = Newton(equations=equations, variable=variable, initial_value=initial_value,
                          e=e, K=K, pot=pot)
    if save==True:
        path = input("请输入保存路径及文件名称：\n")
        try:
            in_one_stage.save_data(path)
            print("保存成功！")
            print("--*--"*10)
        except:
            print("保存失败！请检查相关参数！")
            print("--*--" * 10)
    return in_one_stage.x

if __name__ == '__main__':
    # #例7.2.3 save path:./data/eg_7_2_3
    # one_stage_Newton(2, 'x', start=1, equations=['4*x1-x2+0.1*exp(x1)-1', '-x1+4*x2+(x1**2)/8'], variable=['x1', 'x2'],
    #               initial_value=[0, 0], e=0.001, K=1000)

    # #算例7.3_(1) save path:./data/实例_7_3_1
    # one_stage_Newton(3, 'x', start=1, equations=['x1**2+x2**2+x3**2-1.0', '2*x1**2+x2*2-4*x3', '3*x1**2-4*x2**2+x3**2'],
    #                  variable=['x1', 'x2', 'x3'], initial_value=[1, 1, 1], e=0.001, K=1000)
    #
    #算例7.3_(2) save path:./data/实例_7_3_2
    # one_stage_Newton(2, 'x', start=1, equations=['cos(x1**2+0.4*x2)+x1**2+x2**2-1.6', '1.5*x1**2-1/0.36*x2**2-1'],
    #                  variable=['x1', 'x2'], initial_value=[1.04, 0.47], e=0.001, K=1000, save=False)

    #例7.1.1
    y = []
    x = []
    for i in range(-1200, 1200):
        y.append(one_stage_Newton(1, 'x', start=1, equations=['x1**3-2*x1-5'],
                     variable=['x1'], initial_value=[i], e=0.001, K=1000, save=False))
        x.append(i)
    from matplotlib import pyplot as plt
    plt.figure(2, figsize=(10, 5))
    plt.rcParams["font.sans-serif"] = ["SimHei"]
    plt.rcParams["axes.unicode_minus"] = False
    try:
        plt.scatter(x, y, marker="o", c="c")
        plt.xlabel("迭代初始值")
        plt.ylabel("迭代结果")
    except:
        pass
    plt.show()